/*
 * @Author: huangyupei huangyupei2021@ia.ac.cn
 * @Date: 2024-09-23 14:55:39
 * @LastEditors: huangyupei huangyupei2021@ia.ac.cn
 * @LastEditTime: 2025-03-18 22:13:45
 * @FilePath: /SelfCalib_OptiAcoustic/src/acoustic_BA_g2o.cc
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#include "acoustic_BA_g2o.h"

void SolveAcousticBA(Sophus::SE3d &T_s1_s2, vector<Vector3d> &vPoints, vector<Vector2d> vSonarMeas1, vector<Vector2d> vSonarMeas2, bool priorZ_positive) {
    // pose dimension 6, landmark dimension 3
    // typedef g2o::BlockSolver<g2o::BlockSolverTraits<Eigen::Dynamic, 3>> BlockSolverType;
    // typedef g2o::BlockSolverX BlockSolverType;
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<6, 3>> BlockSolverType;
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;

    // ~ solver 这一块的定义是参考ORB-SLAM3的
    BlockSolverType::LinearSolverType* linearSolver = new LinearSolverType();
    BlockSolverType* solver_ptr = new BlockSolverType(linearSolver);
    auto solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    
    // 定义优化器
    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose(false);        // 是否输出优化过程的详细信息
    
    // 构建 g2o 问题
    // Step 1：添加位姿顶点：固定的第一帧位姿，和待优化的第二帧位姿 T_s1_s2
    g2o::VertexSE3Expmap *vSE3_first_pose = new g2o::VertexSE3Expmap();
    Sophus::SE3d T_s1 = Sophus::SE3d();        // ~ 第一帧的声呐位姿设置为单位阵
    vSE3_first_pose->setEstimate(g2o::SE3Quat(T_s1.unit_quaternion(), T_s1.translation()));
    vSE3_first_pose->setId(0);
    vSE3_first_pose->setFixed(true);          // ~ 第一帧的声呐位姿固定
    optimizer.addVertex(vSE3_first_pose);

    g2o::VertexSE3Expmap *vSE3_second_pose = new g2o::VertexSE3Expmap();
    // Sophus::SE3d T_s1_s2();                // ~ 这里可以想办法把标定整体第一步相机运动估计的结果拿过来当初值    
    vSE3_second_pose->setEstimate(g2o::SE3Quat(T_s1_s2.unit_quaternion(), T_s1_s2.translation()));
    vSE3_second_pose->setId(1);
    vSE3_second_pose->setFixed(false);
    optimizer.addVertex(vSE3_second_pose);

    // Step 2：添加三维点顶点
    const int num_points = 20;
    vector<VertexPoint *> vertex_points;
    for (int i = 0; i < num_points; i++) {
        VertexPoint *v = new VertexPoint();
        v->setId(i + 2);        // 有两个相机位姿顶点，这里保证不与两个位姿顶点的Id重复即可
        v->setEstimate(Vector3d(vPoints[i][0], vPoints[i][1], vPoints[i][2]));  //~ 三维点初值设置为第一帧声呐的观测，表示在第一帧坐标系下
        v->setMarginalized(true);
        optimizer.addVertex(v);
        vertex_points.push_back(v);
    }

    // Step 3：添加声呐观测误差边
    int num_measurements = 20;
    
    // 自由度为2的卡方分布，显著性水平为0.05，对应的临界阈值5.991
    // 可以理解为卡方值高于5.991 95%的几率为外点
    const float deltaMono = sqrt(5.991);
    // const float deltaMono = sqrt(7.815);

    // 第一帧声呐的观测
    for (int i = 0; i < num_measurements; ++i) {
        EdgeSonarProjection *e = new EdgeSonarProjection();
        e->setVertex(0, vSE3_first_pose);
        e->setVertex(1, vertex_points[i]);
        Vector2d first_sonarObs(vSonarMeas1[i][0], vSonarMeas1[i][1]);  //~ 这里的观测是第一帧声呐产生的观测约束
        e->setMeasurement(first_sonarObs);
        e->setInformation(Eigen::Matrix2d::Identity());

        auto rk = new g2o::RobustKernelHuber();
        e->setRobustKernel(rk);
        rk->setDelta(deltaMono);
        // rk->setDelta(1.0);

        optimizer.addEdge(e);
    }

    // 第二帧声呐的观测
    for (int i = 0; i < num_measurements; ++i) {
        EdgeSonarProjection *e = new EdgeSonarProjection();
        e->setVertex(0, vSE3_second_pose);
        e->setVertex(1, vertex_points[i]);
        Vector2d second_sonarObs(vSonarMeas2[i][0], vSonarMeas2[i][1]);    //~ 这里的观测是第二帧声呐产生的观测约束
        e->setMeasurement(second_sonarObs);
        e->setInformation(Matrix2d::Identity());

        auto rk = new g2o::RobustKernelHuber();
        e->setRobustKernel(rk);
        rk->setDelta(deltaMono);

        optimizer.addEdge(e);
    }

    optimizer.initializeOptimization();
    optimizer.optimize(40);
    
    // Step 优化完成后，将想要的结果取出来：第二帧声呐位姿T_s1_s2、重建的三维点
    g2o::VertexSE3Expmap *vSE3_recov = static_cast<g2o::VertexSE3Expmap *>(optimizer.vertex(1));
    auto SE3quat_recov = vSE3_recov->estimate();
    T_s1_s2 = Sophus::SE3d(SE3quat_recov.rotation(), SE3quat_recov.translation());

    for (int i = 0; i < num_points; i++) {
        auto vertex_point = vertex_points[i];
        for (int k = 0; k < 3; k++) vPoints[i][k] = vertex_point->estimate()[k];
    }

    // Step 作为初始化过程，需要利用先验信息确定初始化过程中声呐Z轴运动方向和仰角正负
    double trans_z = T_s1_s2.translation().z();
    // cout << "trans_z of T_s1_s2: " << trans_z << endl;
   
    if ((priorZ_positive && trans_z < 0) || (!priorZ_positive && trans_z > 0)) {
        // 如果知道初始化时Z轴运动方向为正，那么需要保证T_s1_s2的平移部分的Z轴分量为正
        // 如果知道初始化时Z轴运动方向为负，那么需要保证T_s1_s2的平移部分的Z轴分量为负
        // 将T_s1_s2关于镜面对称，然后将vPoints的所有三维点的Z轴取反
        Eigen::Matrix4d P;      // 关于X-Y平面镜面对称
        P << 1, 0, 0, 0,  0, 1, 0, 0,  0, 0, -1, 0,  0, 0, 0, 1;
        
        auto T_dual = P * T_s1_s2.matrix() * P;
        Sophus::SE3d T_s1_s2_dual(T_dual.block<3, 3>(0, 0), T_dual.block<3, 1>(0, 3));
        T_s1_s2 = T_s1_s2_dual;

        for (auto &point: vPoints) {
            point[2] = -point[2];
        }
    }
}

void AcousticPnpBA(Sophus::SE3d &T_s1_si, vector<Vector3d> vPoints, vector<Vector2d> vSonarMeasi) {
    // 构建图优化问题,先设定g2o
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<6, 3>> BlockSolverType;     // pose is 6, landmark is 3
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;    // 线性求解器类型
    
    // 梯度下降方法,可以从GN LM DogLeg选
    BlockSolverType::LinearSolverType* linearSolver = new LinearSolverType();
    BlockSolverType* solver_ptr = new BlockSolverType(linearSolver);
    auto solver = new g2o::OptimizationAlgorithmGaussNewton(solver_ptr);

    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);

    //~ 添加顶点, 只有一个位姿顶点
    g2o::VertexSE3Expmap *vertex_pose = new g2o::VertexSE3Expmap();
    vertex_pose->setEstimate(g2o::SE3Quat(T_s1_si.unit_quaternion(), T_s1_si.translation()));
    vertex_pose->setId(0);
    vertex_pose->setFixed(false);
    optimizer.addVertex(vertex_pose);

    //~ 添加边, 只有3D路标点的重投影观测误差边
    for (int i = 0; i < vSonarMeasi.size(); i++) {
        auto p3d = vPoints[i];
        auto p2d = vSonarMeasi[i];

        EdgePnpSonarProjection *edge = new EdgePnpSonarProjection(p3d);
        edge->setId(i + 1);         // 编号避开顶点编号0即可
        edge->setVertex(0, vertex_pose);
        edge->setMeasurement(p2d);
        edge->setInformation(Eigen::Matrix2d::Identity());
        optimizer.addEdge(edge);
    }

    //~ 开始优化    
    optimizer.setVerbose(false);
    optimizer.initializeOptimization();
    optimizer.optimize(20);

    //~ 优化完成后,取出位姿结果
    g2o::VertexSE3Expmap *vSE3_recov = static_cast<g2o::VertexSE3Expmap *>(optimizer.vertex(0));
    auto SE3quat_recov = vSE3_recov->estimate();
    T_s1_si = Sophus::SE3d(SE3quat_recov.rotation(), SE3quat_recov.translation());
}



void LocalSonarBA(vector<Vector3d> &vPointsSonar, vector<Sophus::SE3d> &vSonarPoses, vector<vector<Vector2d>> vSonarMeasAll) {
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<-1, -1>> BlockSolverType;
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;

    // ~ solver 这一块的定义是参考ORB-SLAM3的
    BlockSolverType::LinearSolverType* linearSolver = new LinearSolverType();
    BlockSolverType* solver_ptr = new BlockSolverType(linearSolver);
    auto solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    
    // 定义优化器
    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose(false);        // 是否输出优化过程的详细信息

    // 优化顶点有两种：初始化计算的3D点 和 第i帧声呐的位姿
    // 构建 g2o 问题
    // Step 1：添加顶点
    vector<VertexPose *> vertex_sonarPoses;
    // 固定顶点只有一个：第一帧声呐的位姿，值为单位阵(实际上可以不必添加这个顶点)
    // 2024-11-27 16:58:31 更新：得添加这个顶点啊！不然少了一个关键的测量信息啊
    VertexPose *vT_s0 = new VertexPose();
    vT_s0->setId(0);
    vT_s0->setEstimate(Sophus::SE3d());
    vT_s0->setFixed(true);
    optimizer.addVertex(vT_s0);
    vertex_sonarPoses.emplace_back(vT_s0);

    // 待优化的位姿顶点：第2~7帧声呐位姿T_si_s0(表示在第一帧声呐坐标系下)
    int num_sonarPose = vSonarPoses.size();     
    for (int i = 1; i < num_sonarPose; i++) {
        VertexPose *vT_si_s0 = new VertexPose();
        vT_si_s0->setId(i);                         // Id范围：1 ~ 6, 对应第2~7帧声呐位姿
        vT_si_s0->setEstimate(vSonarPoses[i].inverse());
        vT_si_s0->setFixed(false);
        optimizer.addVertex(vT_si_s0);
        vertex_sonarPoses.emplace_back(vT_si_s0);       // index范围：1 ~ 6
    }

    // 待优化的3D点顶点：声呐初始化恢复的三维点
    int num_points = vPointsSonar.size();
    vector<VertexPoint *> vertex_points;
    for (int i = 0; i < num_points; i++) {
        VertexPoint *vPoint = new VertexPoint();
        vPoint->setId(i + num_sonarPose);           // 3D点顶点Id从7开始
        vPoint->setEstimate(vPointsSonar[i]);
        vPoint->setMarginalized(false);
        optimizer.addVertex(vPoint);
        vertex_points.emplace_back(vPoint);
    }

    // Step 2: 添加重投影误差边
    // 自由度为2（单相机/声呐，测量为2维）的卡方分布，显著性水平为0.05，对应的临界阈值5.991
    // 可以理解为卡方值高于5.991 95%的几率为外点
    const float deltaMono = sqrt(5.991);
    for (int i = 0; i < num_sonarPose; i++) {      // 从第2幅图像到第7幅图像，总共有6次观测，每次观测有20个误差边【还是要把第1帧带上！！！】
        // 读取声呐图像特征点测量
        vector<Eigen::Vector2d> vSonarMeasi = vSonarMeasAll[i];
        
        // 获取到声呐图像特征点测量后，构建重投影误差边
        for (int j = 0; j < vSonarMeasi.size(); j++) {
            EdgeLocalSonarBAProjection *e = new EdgeLocalSonarBAProjection();
            e->setVertex(0, vertex_sonarPoses[i]);    // 第i幅声呐图像，对应顶点vertex_camPoses[i]
            e->setVertex(1, vertex_points[j]);        // 该误差边观测到的3D点
            e->setMeasurement(vSonarMeasi[j]);
            e->setInformation(Eigen::Matrix2d::Identity());

            auto rk = new g2o::RobustKernelHuber();
            e->setRobustKernel(rk);
            rk->setDelta(deltaMono);

            optimizer.addEdge(e);
        }
    }

    // Step 3: 执行优化
    optimizer.initializeOptimization();
    optimizer.optimize(40);

    // Step 4: 取出优化的结果
    // 取出声呐的位姿
    for (int i = 1; i < num_sonarPose; i++) {
        vSonarPoses[i] = vertex_sonarPoses[i]->estimate().inverse();
    }

    // 取出三维点
    for (int i = 0; i < vPointsSonar.size(); i++) {
        vPointsSonar[i] = vertex_points[i]->estimate();
    }
}